import os

from PIL import Image
"""加黑边变正方形"""
file_dir = '/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/data-test/test_newcut/StraightYellow/'
files = os.listdir(file_dir)
for file in files:
    image = Image.open(file_dir+file)
    w, h = image.size
    target = max(w, h)
    image1 = Image.new("RGB", (target, target))
    if h >= w:
        pad_w = int((target - w) / 2)
        image1.paste(image, (pad_w, 0))
    else:
        pad_h = int((target - h) / 2)
        image1.paste(image, (0, pad_h))
    image1.save(file_dir+file)